{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/pick.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    DiagramBuilder,\n",
    "    MeshcatVisualizer,\n",
    "    Parser,\n",
    "    RigidTransform,\n",
    "    RotationMatrix,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.scenarios import SetColor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DfwIMIJCSkx6"
   },
   "source": [
    "# Compute grasp and pregrasp poses\n",
    "\n",
    "Here is a simple example with a floating Schunk gripper and the foam brick.  It defines the grasp pose as described in the notes, and renders it to the 3D visualizer.\n",
    "\n",
    "**Check yourself**: Try changing the grasp pose to our pregrasp pose.  Do you like the numbers that I picked in the text?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QuFmomDOxoAR"
   },
   "outputs": [],
   "source": [
    "def grasp_poses_example():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "    parser = Parser(plant, scene_graph)\n",
    "    parser.SetAutoRenaming(True)\n",
    "    grasp = parser.AddModelsFromUrl(\n",
    "        \"package://drake/manipulation/models/wsg_50_description/sdf/schunk_wsg_50_no_tip.sdf\"\n",
    "    )[0]\n",
    "    pregrasp = parser.AddModelsFromUrl(\n",
    "        \"package://drake/manipulation/models/wsg_50_description/sdf/schunk_wsg_50_no_tip.sdf\"\n",
    "    )[0]\n",
    "    brick = parser.AddModelsFromUrl(\n",
    "        \"package://drake/examples/manipulation_station/models/061_foam_brick.sdf\"\n",
    "    )[0]\n",
    "    plant.Finalize()\n",
    "\n",
    "    B_O = plant.GetBodyByName(\"base_link\", brick)\n",
    "    B_Ggrasp = plant.GetBodyByName(\"body\", grasp)\n",
    "    B_Gpregrasp = plant.GetBodyByName(\"body\", pregrasp)\n",
    "\n",
    "    # Set the pregrasp to be green and slightly transparent.\n",
    "    inspector = scene_graph.model_inspector()\n",
    "    for body_index in plant.GetBodyIndices(pregrasp):\n",
    "        SetColor(\n",
    "            scene_graph,\n",
    "            [0, 0.6, 0, 0.5],\n",
    "            plant.get_source_id(),\n",
    "            inspector.GetGeometries(plant.GetBodyFrameIdOrThrow(body_index)),\n",
    "        )\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.SetProperty(\"/Background\", \"visible\", False)\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    # TODO(russt): Set a random pose of the object.\n",
    "\n",
    "    # Get the current object, O, pose\n",
    "    X_WO = plant.EvalBodyPoseInWorld(plant_context, B_O)\n",
    "\n",
    "    p_GgraspO = [0, 0.11, 0]\n",
    "    R_GgraspO = RotationMatrix.MakeXRotation(np.pi / 2.0).multiply(\n",
    "        RotationMatrix.MakeZRotation(np.pi / 2.0)\n",
    "    )\n",
    "    X_GgraspO = RigidTransform(R_GgraspO, p_GgraspO)\n",
    "    X_OGgrasp = X_GgraspO.inverse()\n",
    "    X_WGgrasp = X_WO.multiply(X_OGgrasp)\n",
    "\n",
    "    # pregrasp is negative y in the gripper frame (see the figure!).\n",
    "    X_GgraspGpregrasp = RigidTransform([0, -0.08, 0])\n",
    "    X_WGpregrasp = X_WGgrasp @ X_GgraspGpregrasp\n",
    "\n",
    "    plant.SetFreeBodyPose(plant_context, B_Ggrasp, X_WGgrasp)\n",
    "    # Open the fingers, too.\n",
    "    plant.GetJointByName(\"left_finger_sliding_joint\", grasp).set_translation(\n",
    "        plant_context, -0.054\n",
    "    )\n",
    "    plant.GetJointByName(\"right_finger_sliding_joint\", grasp).set_translation(\n",
    "        plant_context, 0.054\n",
    "    )\n",
    "\n",
    "    plant.SetFreeBodyPose(plant_context, B_Gpregrasp, X_WGpregrasp)\n",
    "    # Open the fingers, too.\n",
    "    plant.GetJointByName(\n",
    "        \"left_finger_sliding_joint\", pregrasp\n",
    "    ).set_translation(plant_context, -0.054)\n",
    "    plant.GetJointByName(\n",
    "        \"right_finger_sliding_joint\", pregrasp\n",
    "    ).set_translation(plant_context, 0.054)\n",
    "\n",
    "    diagram.ForcedPublish(context)\n",
    "\n",
    "\n",
    "grasp_poses_example()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Basic Pick and Place.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
